Skip to content

Fix truncation not keeping the last message per the documented rule - #268

Open
Osamaali313 wants to merge 1 commit into
mistralai:mainfrom
Osamaali313:fix/truncate-preserve-last-message
Open

Fix truncation not keeping the last message per the documented rule#268
Osamaali313 wants to merge 1 commit into
mistralai:mainfrom
Osamaali313:fix/truncate-preserve-last-message

Conversation

@Osamaali313

Copy link
Copy Markdown

Problem

_truncate_for_max_tokens (tokens/tokenizers/instruct.py) documents three rules for fitting a conversation into max_tokens, including that the last message must be kept:

# drop some messages to fit in max_tokens. Rules:
# - don't drop any system messages
# - when a user message is dropped, all following assistant|tool message should be dropped until the next
#   user message
# - we never drop the last message

But the inner drop() only guards two cases — system messages and the last user message (last_user_message_index). The last message itself is never guarded. When the final message is an assistant or trailing tool message — finetuning examples always end with the assistant target — the "clear everything until the next user message" cascade can silently exclude it. The result is a truncated example that ends at a user message with no assistant response, plus a Tokenized.prefix_ids that references tokens no longer present in tokens.

This is reachable from the public API: MistralTokenizer.encode_chat_completion(request, max_model_input_len=N) with request.truncate_for_context_length=True.

Fix

Add the missing guard so drop() honors the documented rule:

if idx == len(messages) - 1:
    # never drop the last message
    return

When the last message genuinely cannot fit, to_drop stays positive and the method raises TokenizerException("Input couldn't fit in truncate_at_max_token") — the documented loud failure — instead of returning a corrupted result.

Reproduction

conversation (last msg = final assistant target), tight budget before after
[user, FINAL assistant] last message excluded raises "couldn't fit"
[user, asst, user, FINAL asst] last message excluded raises "couldn't fit"
normal case (a non-final message can be removed) correct correct (unchanged)

Existing test_truncation cases all retain the last message, so they are unaffected; the all-or-nothing case now matches test_truncation_failed's raise.

`_truncate_for_max_tokens` documents that the last message must be kept, but
the inner `drop()` only guards system messages and the last *user* message
(`last_user_message_index`); it never guards the last message itself. When
the final message is an assistant or trailing tool message (finetuning
examples always end with the assistant target), the cascade that clears
everything until the next user message can silently exclude it.

The result is a corrupted example that ends at a user message with no
assistant response, and a `Tokenized.prefix_ids` that references tokens no
longer present. Add the missing guard so the last message is always kept;
when it genuinely cannot fit, `to_drop` stays positive and the method raises
`TokenizerException("Input couldn't fit in truncate_at_max_token")`, the
documented loud failure, instead of returning a corrupted result.

Existing truncation tests always retain the last message, so behavior is
unchanged for them.
Copilot AI review requested due to automatic review settings July 19, 2026 20:34

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants